#!/usr/bin/env python3
import sly
import re
import os
import collections

class PPPException(Exception):
    def __init__(self, msg, pos=None, cause=None):
        if pos is not None:
            super().__init__("Line %d: %s" % (pos.lineno, msg))
        else:
            super().__init__("Line number unknown: %s" % (msg))
class PPPBlockLexer(sly.Lexer):
    tokens = [
        BLOCK_OPEN, BLOCK_SUB, BLOCK_CLOSE, LINE, INLINE_BLOCK
    ]
    reflags = re.M
    @_('\n')
    def ignore_NEWLINE(self, p):
        self.lineno += 1
    BLOCK_SUB = r'^\s*@;.*:\s*?(#.*?)?$'
    BLOCK_OPEN = r'^\s*@\w+.*:\s*?(#.*?)?$'
    BLOCK_CLOSE = r'^\s*@\.\s*?(#.*?)?$'
    #INLINE_BLOCK = r'^\s*@\w+.*\s*?$'
    INLINE_BLOCK = r'^\s*@\w+.*?(#.*?)?$'
    LINE = r'^.+$'

def build_blocks(tokens, path = None):
    stack = []
    stack.append([])
    for tok in tokens:
        line = PPPLine(tok.value, tok.lineno, path)
        if tok.type == 'BLOCK_OPEN':
            stack.append(PPPBlock(line))
        elif tok.type == 'LINE':
            stack[-1].append(line)
        elif tok.type == 'BLOCK_SUB':
            stack[-1].open_sub(line)
        elif tok.type == 'INLINE_BLOCK':
            stack[-1].append(PPPBlock(line))
        elif tok.type == 'BLOCK_CLOSE':
            if type(stack[-1]) == list:
                raise PPPException("Unexcepted BLOCK_CLOSE token", tok)
            stack[-1].close(line)
            cur = stack.pop()
            stack[-1].append(cur)
    ret = stack.pop()
    if type(ret) != list:
        raise PPPException('Unterminated block: "%s"' % ret.cmds[-1].text, ret.cmds[-1])
    return ret

class PPPLineLexer(sly.Lexer):
    tokens = [FREEAT, EVAL_VAR, EVAL_START, STRING1, STRING2, LSQR, RSQR, TEXT]
    EVAL_START = r'@\('
    EVAL_VAR = r'@:\w+'
    FREEAT = r'@'
    STRING1 = r'"([^\\"]|\\.)*"'
    STRING2 = r"'([^\\']|\\.)*'"
    LSQR = r'\('
    RSQR = r'\)'
    TEXT = '[^()\'"@]+'

class PPPLineParser(sly.Parser):
    tokens = PPPLineLexer.tokens
    #debugfile = 'lineparser.out'
    def __init__(self, ctx, mapping=None):
        self.ctx = ctx
        self.mapping = mapping
    @_('group')
    def line(self, p):
        return p[0]
    @_('')
    def line(self, p):
        return ""
    @_('element')
    def group(self, p):
        return p[0]
    @_('group element')
    def group(self, p):
        return p[0] + p[1]
    @_('LSQR RSQR')
    def element(self, p):
        return p[0] + p[1]
    @_('LSQR group RSQR')
    def element(self, p):
        return p[0] + p[1] + p[2]
    @_('EVAL_START group RSQR')
    def element(self, p):
        return str(self.ctx.eval(p[1], mapping=self.mapping))
    @_('EVAL_VAR')
    def element(self, p):
        return str(self.ctx.eval(p[0][2:], mapping=self.mapping))
    @_('FREEAT', 'STRING1', 'STRING2', 'TEXT')
    def element(self, p):
        return p[0]

class PPPBlock:
    def __init__(self, open_cmd):
        self.cmds = [open_cmd]
        self.cur_group = []
        self.groups = [self.cur_group]
    def append(self, child):
        self.cur_group.append(child)
    def open_sub(self, open_cmd):
        self.cmds.append(open_cmd)
        self.cur_group = []
        self.groups.append(self.cur_group)
    def close(self, close_cmd):
        self.close = close_cmd
    def __repr__(self, indent=0):
        ret = ""
        for cmd, group in zip(self.cmds, self.groups):
            ret = ret + cmd.__repr__(indent) + '\n'
            for child in group:
                ret = ret + child.__repr__(indent + 1) + '\n'
        return ret

class PPPLine:
    def __init__(self, text, lineno, path = None):
        self.text = text
        self.lineno = lineno
        self.path = path
    def __repr__(self, indent=0):
        return " " * indent + 'Line %d: %s' % (self.lineno, self.text)

class PPPContext(collections.abc.MutableMapping):
    class Iterator:
        def __init__(self, target):
            self.target = target
            self.curlvl = len(target.stack) - 1
            self.curite = iter(target.stack[self.curlvl])
        def __iter__(self):
            return self
        def __next__(self):
            while True:
                try:
                    nxt = next(self.curite)
                    return nxt
                except StopIteration as e:
                    self.curlvl = self.curlvl - 1
                    if self.curlvl < 0:
                        raise StopIteration('PPPContext iteration stopped')
                    self.curite = iter(self.target.stack[self.curlvl])

    def __init__(self, initial=None):
        self.stack = []
        self.stack.append(initial or {})
        
    def __getitem__(self, key):
        for d in self.stack[::-1]:
            if key in d:
                return d[key]
        raise KeyError('%s not found' % key)
    
    def __setitem__(self, key, v):
        for d in self.stack[::-1]:
            if key in d:
                d[key] = v
        self.stack[-1][key] = v
        
    def __delitem__(self, key, v):
        for d in self.stack[::-1]:
            if key in d:
                del d[key]
                return
        raise KeyError('%s not found' % key)
    
    def __iter__(self):
        return PPPContext.Iterator(self)
    
    def __len__(self):
        total = 0
        for d in self.stack[::-1]:
            total += len(d)
        return total
    
    def push(self, initial = None):
        self.stack.append(initial or {})
        
    def pop(self):
        return self.stack.pop()

    def export(self, key):
        self.stack[-2][key] = self.stack[-1][key]

    def eval(self, st, env = None, mapping = None):
        try:
            return eval(st, env or {}, self)
        except SyntaxError as e:
            lineno = e.lineno - 1
            exc0 = e
        except Exception as e:
            lineno = e.__traceback__.tb_next.tb_lineno - 1
            exc0 = e
        msg = str(exc0)
        if mapping:
            exc = PPPException(msg, mapping[lineno], exc0)
        else:
            exc = PPPException(msg, None, exc0)
        raise exc
    def exec(self, st, env = None, mapping = None):
        try:
            return exec(st, env or {}, self)
        except SyntaxError as e:
            lineno = e.lineno - 1
            exc0 = e
        except Exception as e:
            print(type(e), e)
            lineno = e.__traceback__.tb_next.tb_lineno - 1
            exc0 = e
        print(lineno, mapping, exc0)
        msg = str(exc0)
        if mapping:
            exc = PPPException(msg, mapping[lineno], exc0)
        else:
            exc = PPPException(msg, None, exc0)
        raise exc
        return 

expand_text = lambda text, ctx, mapping=None: PPPLineParser(ctx, mapping).parse(PPPLineLexer().tokenize(text))
expand_line = lambda line, ctx: expand_text(line.text, ctx, [line])

CMD0_RE = re.compile(r'\s*@;?(?P<cmd>\w+)\s*(?P<arg>.*):\s*(#.*?)?$')
CMD1_RE = re.compile(r'\s*@(?P<cmd>\w+)\s*(?P<arg>.*?)\s*(#.*?)?$')
match_cmd = lambda cmd: CMD0_RE.match(cmd.text) or CMD1_RE.match(cmd.text)

def prepare_cmd(cmd, ctx, expand_args = True):
    m = match_cmd(cmd)
    if m:
        if expand_args:
            return m['cmd'], (expand_text(m['arg'], ctx) or "").strip()
        else:
            return m['cmd'], m['arg'].strip()
    raise PPPException('Command format error.', cmd)

def run_for(block, ctx, blockmap):
    cmd, arg = prepare_cmd(block.cmds[0], ctx)
    ret = []
    run_child = lambda: ret.extend(run_group(block.groups[0], ctx, blockmap))
    ctx.exec('for %s: run_child()' % arg, {'run_child': run_child}, mapping=block.cmds)
    return ret

def run_if(block, ctx, blockmap):
    ret = []
    def run_child(i):
        ret.extend(run_group(block.groups[i], ctx, blockmap))
    pycmds = []
    for i, cmd in enumerate(block.cmds):
        cmd, arg = prepare_cmd(cmd, ctx)
        pycmds.append("%s %s: run_child(%d)" % (cmd, arg, i))
    pycmd = "\n".join(pycmds)
    ctx.exec(pycmd, {'run_child': run_child}, mapping=block.cmds)
    return ret

WS_RE = re.compile(r'\s+')
def run_define(block, ctx, blockmap):
    cmd, arg = prepare_cmd(block.cmds[0], ctx, False)
    name, value = WS_RE.split(arg, maxsplit=1)
    ctx[name] = value
    return []

def run_set(block, ctx, blockmap):
    cmd, arg = prepare_cmd(block.cmds[0], ctx)
    name, value = WS_RE.split(arg, maxsplit=1)
    ctx[name] = ctx.eval(value, mapping=block.cmds)
    return []

def run_macro(block, ctx, blockmap):
    cmd, arg = prepare_cmd(block.cmds[0], ctx)
    ctx.push()
    ctx.exec("def %s: return locals()" % arg, mapping=block.cmds)
    name, get_args = next(iter(ctx.pop().items()))
    group = block.groups[0]
    defcmd = block.cmds[0]
    name = get_args.__name__
    def expand(block, ctx, blockmap):
        ret = []
        cmd, args = prepare_cmd(block.cmds[0], ctx)

        formal_args = ctx.eval('__get_args__%s' % args, {'__get_args__': get_args}, mapping = block.cmds)
        ctx.push(formal_args)
        run_child = lambda: ret.extend(run_group(group, ctx, blockmap))
        ctx.exec('run_child()', {'run_child': run_child}, mapping=block.cmds)
        ctx.pop()
        return ret
    blockmap[name] = expand
    return []

def run_env(block, ctx, blockmap):
    cmd, arg = prepare_cmd(block.cmds[0], ctx)
    ctx.push()
    ctx.exec("def %s: return locals()" % arg, mapping=block.cmds)
    name, get_args = next(iter(ctx.pop().items()))
    group = block.groups[0]    
    name = get_args.__name__
    defcmd = block.cmds[0]
    def expand(block, ctx, blockmap):
        ret = []
        cmd, args = prepare_cmd(block.cmds[0], ctx)
        formal_args = ctx.eval('__get_args__%s' % args, {'__get_args__': get_args}, mapping=block.cmds)
        ctx.push(formal_args)
        groupmap = {"" : (block.cmds[0], block.groups[0])}
        for rawcmd, rawgroup in zip(block.cmds[1:], block.groups[1:]):
            cmd, args = prepare_cmd(rawcmd, ctx)
            groupmap[cmd] = (rawcmd, rawgroup)
        def paste(block, ctx, blockmap):
            cmd, args = prepare_cmd(block.cmds[0], ctx)
            ret = []
            ret.extend(run_group(groupmap[args][1], ctx, blockmap))
            return ret
        blockmap.push({'paste': paste})
        run_child = lambda: ret.extend(run_group(group, ctx, blockmap))
        ctx.exec('run_child()', {'run_child': run_child}, mapping=block.cmds)
        ctx.pop()
        blockmap.pop()
        return ret
    blockmap[name] = expand
    return []

def run_py(block, ctx, blockmap):
    code = []
    for line in block.groups[0]:
        code.append(line.text)
    ctx.exec("\n".join(code), mapping=block.groups[0])
    return []

def run_include(block, ctx, blockmap):
    cmd, arg = prepare_cmd(block.cmds[0], ctx)
    srcdir = os.path.dirname(block.cmds[0].path)
    incfile = None
    if os.path.exists(os.path.join(srcdir, arg)):
        incfile = os.path.join(srcdir, arg)
    elif 'incpath' in ctx:
        for incdir in ctx['incpath']:
            if os.path.exists(os.path.join(incdir, arg)):
                incfile = os.path.join(incdir, arg)
    if incfile is None:
        raise PPPException('Cannot find file "%s"' % arg, block.cmds[0])
    return run_file(incfile, ctx, blockmap)

def run_push(block, ctx, blockmap):
    ctx.push({})
    return []
def run_pop(block, ctx, blockmap):
    ctx.pop()
    return []
SP_RE=re.compile(r'\s*,\s*')
def run_export(block, ctx, blockmap):
    cmd, args = prepare_cmd(block.cmds)
    for arg in SP_RE.split(args):
        ctx.export(arg)
    return []
def run_file(path, ctx = None, blockmap = None):
    text = open(path).read()
    lexer = PPPBlockLexer()
    blks = build_blocks(lexer.tokenize(text), path)
    ret = []
    ctx = ctx if ctx is not None else PPPContext()
    blockmap = blockmap or PPPContext(dict(blockmap_default))
    ret = run_group(blks, ctx, blockmap)
    return ret

def run_group(group, ctx, blockmap):
    ret = []
    for child in group:
        if type(child) == PPPLine:
            line_exp = expand_line(child, ctx)
            ret.append((line_exp, child.path, child.lineno))
        else:
            ret.extend(run_block(child, ctx, blockmap))
    return ret

def concat_results(lines, linenote):
    ret = []
    path_last = None
    line_last = -10
    for text, path, lineno in lines:
        if lineno != line_last + 1 or path != path_last:
            ret.append(linenote(path, lineno))
        ret.append(text)
        line_last = lineno
        path_last = path
    return "\n".join(ret)

blockmap_default = {
    "for": run_for,
    "if": run_if,
    "define": run_define,
    "set": run_set,
    "macro": run_macro,
    "env": run_env,
    "py": run_py,
    "include": run_include,
    "push": run_push,
    "pop": run_pop,
    "export": run_export
}

PERMISSIVE_RUN_BLOCK = False
def run_block(block, ctx, blockmap = blockmap_default, linenote=lambda line: ""):
    try:
        ret = []
        cmd = match_cmd(block.cmds[0])['cmd']
        if cmd in blockmap:
            return blockmap[cmd](block, ctx, blockmap)
        if not PERMISSIVE_RUN_BLOCK:
            raise PPPException("Unknown block %s" % cmd, block.cmds[0])
        for group in block.groups:
            ret.extend(run_group(group, ctx, blockmap))
        return ret
    except:
        raise 
def clinenote(line, delta):
    return '#line %d "%s"' % (line.lineno + delta, line.path)
import sys
import json
if __name__ == '__main__':
    ret = run_file(sys.argv[1])
    print(concat_results(ret, lambda path, line: '#line %d "%s"' % (line, path)))
